#ifndef     __RT_TASK_AWAITER2_H
#define     __RT_TASK_AWAITER2_H

#include <coroutine>
#include <type_traits>

namespace rt {

template <typename PROMISE, typename T>
class TaskAwaiter2 {

public:
    TaskAwaiter2(std::coroutine_handle<PROMISE> handle) : m_handle(handle) {
    }

    bool await_ready() const noexcept {
        return m_handle.done();
    }

    void await_suspend(std::coroutine_handle<> handle) {
        PROMISE& promise = m_handle.promise();
        promise.set_parent(handle);
    }

    T await_resume() {
        PROMISE& promise = m_handle.promise();
        if constexpr (std::is_void_v<T>) {
            promise.get();
        } else {
            return promise.get();
        }
    }

private:
    std::coroutine_handle<PROMISE> m_handle;
};

}

#endif
